# coding=utf-8
# Copyright 2019 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Run a classifier on MNIST and test OOD."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import app
from absl import flags
import tensorflow as tf  # tf
import tensorflow_probability as tfp
from extrapolation.classifier.classifier import CNN as CNN
import extrapolation.classifier.train_cnn as train_cnn
from extrapolation.utils import dataset_utils
from extrapolation.utils import utils

tfd = tfp.distributions


flags.DEFINE_string('ood_classes', '5', 'a comma-separated list of'
                    'which labels to consider OoD')

FLAGS = flags.FLAGS



def main(argv):
  if len(argv) > 1:
    raise app.UsageError('Too many command-line arguments.')

  params = FLAGS.flag_values_dict()
  tf.set_random_seed(params['seed'])

  params['results_dir'] = utils.make_subdir(
      params['results_dir'], params['expname'])
  params['figdir'] = utils.make_subdir(params['results_dir'], 'figs')
  params['ckptdir'] = utils.make_subdir(params['results_dir'], 'ckpts')
  params['logdir'] = utils.make_subdir(params['results_dir'], 'logs')
  params['tensordir'] = utils.make_subdir(params['results_dir'], 'tensors')

  ood_classes = [int(x) for x in params['ood_classes'].split(',')]
  # We assume we train on all non-OOD classes.
  all_classes = range(params['n_classes'])
  ind_classes = [x for x in all_classes if x not in ood_classes]
  (itr_train,
   itr_valid,
   itr_test,
   itr_test_ood) = dataset_utils.load_dset_ood_supervised(ind_classes,
                                                          ood_classes)

  conv_dims = [int(x) for x in params['conv_dims'].split(',')]
  conv_sizes = [int(x) for x in params['conv_sizes'].split(',')]
  clf = CNN(conv_dims, conv_sizes, params['n_classes'])
  params['n_layers'] = len(conv_dims)

  train_cnn.train_classifier(clf, itr_train, itr_valid, params)
  train_cnn.test_classifier(clf, itr_test, params, 'test')

  params['tensordir'] = utils.make_subdir(
      params['results_dir'], 'train_tensors')
  train_cnn.test_classifier(clf, itr_train, params, 'train')

  params['tensordir'] = utils.make_subdir(
      params['results_dir'], 'ood_tensors')
  train_cnn.test_classifier(clf, itr_test_ood, params, 'ood')

if __name__ == '__main__':
  app.run(main)
